{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b698676e-9a49-489d-a396-d717a4ef6b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindspore import context\n",
    "from mindsponge import Sponge\n",
    "from mindsponge import Molecule\n",
    "from mindsponge import ForceField\n",
    "from mindsponge import UpdaterMD\n",
    "from mindsponge.callback import WriteH5MD, RunInfo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "71c58271-41cd-4f15-a88a-4ac371a09a61",
   "metadata": {},
   "outputs": [],
   "source": [
    "context.set_context(mode=context.GRAPH_MODE, device_target=\"GPU\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "103887d0-1c8b-49fe-933c-27cec1001c3f",
   "metadata": {},
   "source": [
    "## 使用力场模板\n",
    "\n",
    "MindSPONGE中支持yaml格式的力场参数信息配置，其中大体分为`模板`和`参数`两个大类。`模板`中主要存储系统中的原子名称、原子类型、原子质量、原子序数、成键关系等基本信息，也是一般都不需要去改动的信息。`参数`中顾名思义就是存储一些力场参数，比如键长相互作用$E_b=k_b(l-l_0)^2$中的$(k_b,l_0)$就是一组力场参数。不同的原子组合形成的键、角度、二面角还有其他的一些远程相互作用等，都会对应于不同的参数，在不同的力场中千差万别。\n",
    "\n",
    "比如这是一个三点水的模板配置信息：\n",
    "\n",
    "```yaml\n",
    "template:\n",
    "  WAT:\n",
    "    atom_name: [O, H1, H2]\n",
    "    atom_type: [OW, HW, HW]\n",
    "    atom_mass: [16.00, 1.008, 1.008]\n",
    "    atomic_number: [8, 1, 1]\n",
    "    bond:\n",
    "    - [0, 1]\n",
    "    - [0, 2]\n",
    "    head_atom: null\n",
    "    tail_atom: null\n",
    "```\n",
    "\n",
    "以及一个三点水的力场参数信息：\n",
    "\n",
    "```yaml\n",
    "template:\n",
    "  base: water.tip3p.yaml\n",
    "parameters:\n",
    "  bond_energy:\n",
    "    length_unit: nm\n",
    "    energy_unit: kj/mol\n",
    "    parameter_names:\n",
    "      atoms: [bond_length, force_constant]\n",
    "    parameters:\n",
    "      OW-HW: [0.09572, 502416]\n",
    "  angle_energy:\n",
    "    length_unit: nm\n",
    "    energy_unit: kj/mol\n",
    "    parameter_names:\n",
    "      atoms: [bond_angle, force_constant]\n",
    "    parameters:\n",
    "      HW-OW-HW: [1.824218, 628.02]\n",
    "  coulomb_energy:\n",
    "    length_unit: nm\n",
    "    energy_unit: kj/mol\n",
    "  vdw_energy:\n",
    "    length_unit: nm\n",
    "    energy_unit: kj/mol\n",
    "    parameter_names:\n",
    "      atoms: [sigma, epsilon]\n",
    "    parameters:\n",
    "      OW: [0.315061, 0.636386]\n",
    "      HW: [0.0, 0.0]\n",
    "```\n",
    "\n",
    "用户可以根据不同的力场自定义一个这样的yaml文件，然后就可以在MindSPONGE平台上非常方便的进行分子动力学模拟。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f4b59501-20e9-4b79-a33d-0376ef25232d",
   "metadata": {},
   "outputs": [],
   "source": [
    "system = Molecule(template='water.tip3p.yaml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "12adc305-1db8-47a8-bb6f-df1ad83a2040",
   "metadata": {},
   "outputs": [],
   "source": [
    "potential = ForceField(system, parameters='TIP3P')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "56f70d3b-92e3-41e4-8119-ef0395e8b279",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = UpdaterMD(system=system,\n",
    "                time_step=1e-3,\n",
    "                integrator='velocity_verlet',\n",
    "                temperature=300,\n",
    "                thermostat='langevin',)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1fb210d5-8c8d-4f79-848f-d4cd3d5549a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MindSPONGE] Using fully connected neighbour list (not updated).\n"
     ]
    }
   ],
   "source": [
    "md = Sponge(system, potential, opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9e28a60d-fe84-4afb-8c94-37fda8b56bba",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_info = RunInfo(10)\n",
    "cb_h5md = WriteH5MD(system, 'tutorial_b02.h5md', save_freq=10, write_velocity=True, write_force=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "26add33b-e30b-4940-9aa7-8231d56e96f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MindSPONGE] Started simulation at 2023-06-30 16:23:30\n",
      "[MindSPONGE] Step: 0, E_pot: 9.203461, E_kin: 0.0, E_tot: 9.203461, Temperature: 0.0\n",
      "[MindSPONGE] Step: 10, E_pot: 1.6198775, E_kin: 3.9923854, E_tot: 5.6122627, Temperature: 160.05788\n",
      "[MindSPONGE] Step: 20, E_pot: 0.57855374, E_kin: 5.447758, E_tot: 6.026312, Temperature: 218.4049\n",
      "[MindSPONGE] Step: 30, E_pot: 1.6316944, E_kin: 2.4541254, E_tot: 4.0858197, Temperature: 98.38782\n",
      "[MindSPONGE] Step: 40, E_pot: 1.5366576, E_kin: 2.3814979, E_tot: 3.9181554, Temperature: 95.47612\n",
      "[MindSPONGE] Step: 50, E_pot: 0.73025894, E_kin: 2.880464, E_tot: 3.610723, Temperature: 115.48007\n",
      "[MindSPONGE] Step: 60, E_pot: 1.8875588, E_kin: 4.9203844, E_tot: 6.8079433, Temperature: 197.26208\n",
      "[MindSPONGE] Step: 70, E_pot: 1.8393941, E_kin: 4.995426, E_tot: 6.8348203, Temperature: 200.27057\n",
      "[MindSPONGE] Step: 80, E_pot: 1.5443127, E_kin: 7.7986975, E_tot: 9.34301, Temperature: 312.6559\n",
      "[MindSPONGE] Step: 90, E_pot: 1.0576841, E_kin: 7.809618, E_tot: 8.867302, Temperature: 313.09372\n",
      "[MindSPONGE] Step: 100, E_pot: 1.832165, E_kin: 5.801772, E_tot: 7.633937, Temperature: 232.59761\n",
      "[MindSPONGE] Step: 110, E_pot: 1.1987792, E_kin: 4.1811476, E_tot: 5.3799267, Temperature: 167.62549\n",
      "[MindSPONGE] Step: 120, E_pot: 0.8765099, E_kin: 5.436157, E_tot: 6.312667, Temperature: 217.93982\n",
      "[MindSPONGE] Step: 130, E_pot: 3.1524642, E_kin: 8.614393, E_tot: 11.766857, Temperature: 345.35782\n",
      "[MindSPONGE] Step: 140, E_pot: 1.0800122, E_kin: 10.348159, E_tot: 11.428171, Temperature: 414.86584\n",
      "[MindSPONGE] Step: 150, E_pot: 0.8408035, E_kin: 7.781229, E_tot: 8.622032, Temperature: 311.9556\n",
      "[MindSPONGE] Step: 160, E_pot: 4.7319903, E_kin: 4.7583647, E_tot: 9.490355, Temperature: 190.76659\n",
      "[MindSPONGE] Step: 170, E_pot: 0.07747452, E_kin: 7.1626964, E_tot: 7.240171, Temperature: 287.1581\n",
      "[MindSPONGE] Step: 180, E_pot: 3.6156647, E_kin: 9.619015, E_tot: 13.234679, Temperature: 385.63385\n",
      "[MindSPONGE] Step: 190, E_pot: 2.2509513, E_kin: 12.884341, E_tot: 15.135292, Temperature: 516.54333\n",
      "[MindSPONGE] Step: 200, E_pot: 2.371033, E_kin: 19.366741, E_tot: 21.737774, Temperature: 776.42786\n",
      "[MindSPONGE] Step: 210, E_pot: 1.9925103, E_kin: 12.139831, E_tot: 14.132341, Temperature: 486.69537\n",
      "[MindSPONGE] Step: 220, E_pot: 0.94162667, E_kin: 17.088982, E_tot: 18.03061, Temperature: 685.1107\n",
      "[MindSPONGE] Step: 230, E_pot: 3.583885, E_kin: 13.051564, E_tot: 16.635448, Temperature: 523.2475\n",
      "[MindSPONGE] Step: 240, E_pot: 1.8338804, E_kin: 16.327486, E_tot: 18.161366, Temperature: 654.5817\n",
      "[MindSPONGE] Step: 250, E_pot: 2.1623373, E_kin: 18.10532, E_tot: 20.267658, Temperature: 725.8565\n",
      "[MindSPONGE] Step: 260, E_pot: 0.9804783, E_kin: 22.500973, E_tot: 23.481451, Temperature: 902.08167\n",
      "[MindSPONGE] Step: 270, E_pot: 3.806062, E_kin: 20.907413, E_tot: 24.713476, Temperature: 838.19464\n",
      "[MindSPONGE] Step: 280, E_pot: 4.0621953, E_kin: 16.893906, E_tot: 20.9561, Temperature: 677.29\n",
      "[MindSPONGE] Step: 290, E_pot: 0.48108497, E_kin: 21.001957, E_tot: 21.483042, Temperature: 841.985\n",
      "[MindSPONGE] Step: 300, E_pot: 2.2917974, E_kin: 15.985594, E_tot: 18.277391, Temperature: 640.875\n",
      "[MindSPONGE] Step: 310, E_pot: 3.7466593, E_kin: 17.236282, E_tot: 20.98294, Temperature: 691.0161\n",
      "[MindSPONGE] Step: 320, E_pot: 8.304243, E_kin: 10.076683, E_tot: 18.380926, Temperature: 403.98215\n",
      "[MindSPONGE] Step: 330, E_pot: 2.218204, E_kin: 6.8653874, E_tot: 9.083591, Temperature: 275.23877\n",
      "[MindSPONGE] Step: 340, E_pot: 2.3807073, E_kin: 7.6315084, E_tot: 10.012216, Temperature: 305.95316\n",
      "[MindSPONGE] Step: 350, E_pot: 1.7648168, E_kin: 5.5101933, E_tot: 7.27501, Temperature: 220.90797\n",
      "[MindSPONGE] Step: 360, E_pot: 2.3034525, E_kin: 3.7670364, E_tot: 6.070489, Temperature: 151.02345\n",
      "[MindSPONGE] Step: 370, E_pot: 1.0802875, E_kin: 6.515034, E_tot: 7.5953217, Temperature: 261.19284\n",
      "[MindSPONGE] Step: 380, E_pot: 0.36301306, E_kin: 9.713495, E_tot: 10.0765085, Temperature: 389.42166\n",
      "[MindSPONGE] Step: 390, E_pot: 5.9782686, E_kin: 8.318791, E_tot: 14.29706, Temperature: 333.50687\n",
      "[MindSPONGE] Step: 400, E_pot: 4.892672, E_kin: 5.505479, E_tot: 10.398151, Temperature: 220.71898\n",
      "[MindSPONGE] Step: 410, E_pot: 5.5998273, E_kin: 5.9521704, E_tot: 11.551998, Temperature: 238.6272\n",
      "[MindSPONGE] Step: 420, E_pot: 4.9029346, E_kin: 5.576167, E_tot: 10.479101, Temperature: 223.55292\n",
      "[MindSPONGE] Step: 430, E_pot: 0.7374965, E_kin: 6.5696645, E_tot: 7.307161, Temperature: 263.38303\n",
      "[MindSPONGE] Step: 440, E_pot: 3.9987838, E_kin: 2.2944741, E_tot: 6.2932577, Temperature: 91.98727\n",
      "[MindSPONGE] Step: 450, E_pot: 2.8738399, E_kin: 1.9683135, E_tot: 4.8421535, Temperature: 78.91124\n",
      "[MindSPONGE] Step: 460, E_pot: 1.0607338, E_kin: 7.2753553, E_tot: 8.336089, Temperature: 291.6747\n",
      "[MindSPONGE] Step: 470, E_pot: 0.86768323, E_kin: 7.432702, E_tot: 8.300385, Temperature: 297.98288\n",
      "[MindSPONGE] Step: 480, E_pot: 0.8224413, E_kin: 5.0779123, E_tot: 5.9003534, Temperature: 203.5775\n",
      "[MindSPONGE] Step: 490, E_pot: 2.1243675, E_kin: 5.1672583, E_tot: 7.291626, Temperature: 207.15944\n",
      "[MindSPONGE] Step: 500, E_pot: 2.8870032, E_kin: 6.974619, E_tot: 9.861622, Temperature: 279.61795\n",
      "[MindSPONGE] Step: 510, E_pot: 0.85057014, E_kin: 11.598957, E_tot: 12.449527, Temperature: 465.0113\n",
      "[MindSPONGE] Step: 520, E_pot: 3.9154315, E_kin: 8.942133, E_tot: 12.857565, Temperature: 358.49713\n",
      "[MindSPONGE] Step: 530, E_pot: 4.489691, E_kin: 10.046013, E_tot: 14.535704, Temperature: 402.75256\n",
      "[MindSPONGE] Step: 540, E_pot: 5.9207582, E_kin: 10.768467, E_tot: 16.689224, Temperature: 431.7163\n",
      "[MindSPONGE] Step: 550, E_pot: 2.8750532, E_kin: 11.440697, E_tot: 14.31575, Temperature: 458.6665\n",
      "[MindSPONGE] Step: 560, E_pot: 1.169558, E_kin: 10.325634, E_tot: 11.495192, Temperature: 413.9628\n",
      "[MindSPONGE] Step: 570, E_pot: 5.902687, E_kin: 6.181387, E_tot: 12.084074, Temperature: 247.81665\n",
      "[MindSPONGE] Step: 580, E_pot: 2.9332042, E_kin: 11.212723, E_tot: 14.145927, Temperature: 449.5269\n",
      "[MindSPONGE] Step: 590, E_pot: 0.32937884, E_kin: 17.637583, E_tot: 17.96696, Temperature: 707.10455\n",
      "[MindSPONGE] Step: 600, E_pot: 1.2222682, E_kin: 14.300676, E_tot: 15.522944, Temperature: 573.3253\n",
      "[MindSPONGE] Step: 610, E_pot: 2.6599123, E_kin: 10.857535, E_tot: 13.517447, Temperature: 435.28714\n",
      "[MindSPONGE] Step: 620, E_pot: 3.1567626, E_kin: 9.646782, E_tot: 12.803545, Temperature: 386.74707\n",
      "[MindSPONGE] Step: 630, E_pot: 2.6973917, E_kin: 9.057779, E_tot: 11.755171, Temperature: 363.13348\n",
      "[MindSPONGE] Step: 640, E_pot: 0.24941514, E_kin: 13.235463, E_tot: 13.484879, Temperature: 530.6201\n",
      "[MindSPONGE] Step: 650, E_pot: 2.1370502, E_kin: 9.524891, E_tot: 11.661942, Temperature: 381.86038\n",
      "[MindSPONGE] Step: 660, E_pot: 3.9134638, E_kin: 8.236693, E_tot: 12.150157, Temperature: 330.2155\n",
      "[MindSPONGE] Step: 670, E_pot: 4.2805233, E_kin: 5.052934, E_tot: 9.333458, Temperature: 202.57611\n",
      "[MindSPONGE] Step: 680, E_pot: 3.3028224, E_kin: 6.2983184, E_tot: 9.601141, Temperature: 252.50453\n",
      "[MindSPONGE] Step: 690, E_pot: 0.7145305, E_kin: 13.880701, E_tot: 14.595232, Temperature: 556.4882\n",
      "[MindSPONGE] Step: 700, E_pot: 5.2229376, E_kin: 13.072342, E_tot: 18.29528, Temperature: 524.08044\n",
      "[MindSPONGE] Step: 710, E_pot: 5.4436226, E_kin: 9.829201, E_tot: 15.272823, Temperature: 394.0604\n",
      "[MindSPONGE] Step: 720, E_pot: 4.220007, E_kin: 7.6816635, E_tot: 11.90167, Temperature: 307.96393\n",
      "[MindSPONGE] Step: 730, E_pot: 3.767537, E_kin: 10.086422, E_tot: 13.853959, Temperature: 404.3726\n",
      "[MindSPONGE] Step: 740, E_pot: 3.7170727, E_kin: 10.590387, E_tot: 14.30746, Temperature: 424.57697\n",
      "[MindSPONGE] Step: 750, E_pot: 10.167873, E_kin: 7.57446, E_tot: 17.742332, Temperature: 303.66608\n",
      "[MindSPONGE] Step: 760, E_pot: 3.817353, E_kin: 15.734091, E_tot: 19.551443, Temperature: 630.7921\n",
      "[MindSPONGE] Step: 770, E_pot: 4.2830634, E_kin: 20.356342, E_tot: 24.639406, Temperature: 816.1018\n",
      "[MindSPONGE] Step: 780, E_pot: 5.547362, E_kin: 16.67034, E_tot: 22.217701, Temperature: 668.327\n",
      "[MindSPONGE] Step: 790, E_pot: 6.798376, E_kin: 13.479874, E_tot: 20.27825, Temperature: 540.41876\n",
      "[MindSPONGE] Step: 800, E_pot: 19.116703, E_kin: 8.12423, E_tot: 27.240932, Temperature: 325.7068\n",
      "[MindSPONGE] Step: 810, E_pot: 12.37052, E_kin: 6.828092, E_tot: 19.198612, Temperature: 273.7436\n",
      "[MindSPONGE] Step: 820, E_pot: 4.272229, E_kin: 17.572552, E_tot: 21.84478, Temperature: 704.49744\n",
      "[MindSPONGE] Step: 830, E_pot: 9.549255, E_kin: 14.505956, E_tot: 24.05521, Temperature: 581.5552\n",
      "[MindSPONGE] Step: 840, E_pot: 12.38779, E_kin: 12.302016, E_tot: 24.689806, Temperature: 493.1975\n",
      "[MindSPONGE] Step: 850, E_pot: 11.626081, E_kin: 9.875748, E_tot: 21.50183, Temperature: 395.92648\n",
      "[MindSPONGE] Step: 860, E_pot: 6.34471, E_kin: 10.602924, E_tot: 16.947634, Temperature: 425.0796\n",
      "[MindSPONGE] Step: 870, E_pot: 7.9197545, E_kin: 8.942096, E_tot: 16.86185, Temperature: 358.49564\n",
      "[MindSPONGE] Step: 880, E_pot: 3.9619708, E_kin: 8.946284, E_tot: 12.908255, Temperature: 358.66357\n",
      "[MindSPONGE] Step: 890, E_pot: 5.861244, E_kin: 11.982002, E_tot: 17.843246, Temperature: 480.3679\n",
      "[MindSPONGE] Step: 900, E_pot: 7.0596886, E_kin: 5.1685286, E_tot: 12.228217, Temperature: 207.21037\n",
      "[MindSPONGE] Step: 910, E_pot: 3.4878178, E_kin: 13.520969, E_tot: 17.008787, Temperature: 542.0663\n",
      "[MindSPONGE] Step: 920, E_pot: 6.8617764, E_kin: 10.557119, E_tot: 17.418896, Temperature: 423.24323\n",
      "[MindSPONGE] Step: 930, E_pot: 4.2134385, E_kin: 8.868448, E_tot: 13.081886, Temperature: 355.54306\n",
      "[MindSPONGE] Step: 940, E_pot: 1.0145078, E_kin: 10.65641, E_tot: 11.6709175, Temperature: 427.22385\n",
      "[MindSPONGE] Step: 950, E_pot: 0.085960776, E_kin: 6.46329, E_tot: 6.549251, Temperature: 259.1184\n",
      "[MindSPONGE] Step: 960, E_pot: 4.432596, E_kin: 5.3544345, E_tot: 9.787031, Temperature: 214.66348\n",
      "[MindSPONGE] Step: 970, E_pot: 2.8022656, E_kin: 8.253837, E_tot: 11.056103, Temperature: 330.9028\n",
      "[MindSPONGE] Step: 980, E_pot: 4.1750045, E_kin: 7.592721, E_tot: 11.767725, Temperature: 304.39813\n",
      "[MindSPONGE] Step: 990, E_pot: 4.8383408, E_kin: 12.084518, E_tot: 16.92286, Temperature: 484.47787\n",
      "[MindSPONGE] Finished simulation at 2023-06-30 16:23:41\n",
      "[MindSPONGE] Simulation time: 10.25 seconds.\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<mindsponge.core.sponge.Sponge at 0x7f1e6cc02070>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "md.run(1000, callbacks=[run_info, cb_h5md])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "89a5690f-75e9-4ca4-b133-5e50932f9101",
   "metadata": {},
   "source": [
    "![](../../docs/tutorial_b02.gif)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mindspore-1.9",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "768f82fbbece3f0acd424cda825de8aa44edbfaa38d8a3df4ff1437f7ae7861f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
